// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: daniel

#pragma once

#include <memory>
#include <string>
#include <unordered_map>
#include <utility>

#include <nlohmann/json.hpp>

#include "coin2/exchange/base/order/rate_tracker.h"
#include "coin2/exchange/order/connection.pb.h"

constexpr auto EXAMPLE_SPOT_KRAKEN = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke06/kraken/trade_key.json",
  "rate_limit_info": {
    "spot_kraken": {
      "req_all": {
        "secs": 48,
        "reqs": 180
      }
    }
  }
})";

class KrakenRateTracker : public coin2::exchange::base::order::IRateTracker {
 public:
  using RateTracker = coin2::exchange::base::order::RateTracker;
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using NativeClientConfig = coin2::exchange::order::OrderConnectionConfig;
  using RateLimitReportWriter = coin2::exchange::base::order::RateLimitReportWriter;

 public:
  explicit KrakenRateTracker(const MarketExchangeApi& mea, const NativeClientConfig& config)
      : rate_limit_report_writer_{nullptr} {
    using RateConfig = coin2::exchange::base::order::RateConfig;
    coin2::exchange::order::RateLimitInfoSpotKraken rate_config;
    RateConfig::Get()->GetConfig(mea, config, &rate_config);
    if (!rate_config.has_req_all()) {
      RateConfig::CheckFail(mea.String(), EXAMPLE_SPOT_KRAKEN);
    }
    const auto& req_all = rate_config.req_all();
    LOG(INFO) << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
              << "\nOG request rate was set as below:\n"
              << "\nreq_all: [" << req_all.reqs() << " (reqs) / " << req_all.secs() << " (seconds)]"
              << "\n\nOr you can configure customized rate in driver config. For example: "
              << EXAMPLE_SPOT_KRAKEN
              << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";
    rate_limit_info_ = rate_config;

    // rate limit report
    rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
    rate_limit_report_writer_->SetRateLimitInfoString(&rate_config);
  }

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) override { CHECK(false); }

  bool HasEnoughRateLimit(int64_t timestamp, ActionType action_type) override { CHECK(false); }

  void Push(int32_t num_reqs, int64_t timestamp, [[maybe_unused]] ActionType action_type) override {
    CHECK(false);
  }

  void ParseRateLimitInfo(
      [[maybe_unused]] const std::shared_ptr<AsioHttpContext>& http_context,
      [[maybe_unused]] ActionType action_type) override {}

  void SetRateLimitInfoString(const google::protobuf::Message* message) override {
    rate_limit_report_writer_->SetRateLimitInfoString(message);
  }

  void WriteRateLimitReportMessage(OrderLoggerCB callback) override {
    if (!callback) {
      return;
    }
    rate_limit_report_writer_->Write(callback);
  }

  double GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) override {
    // TODO(daniel):
    return 0.;
  }

  bool HasEnoughRateLimit(int64_t ts, ActionType action_type, const std::string& symbol) {
    bool ret = true;
    switch (action_type) {
      case ActionType::PLACE:
      case ActionType::AMEND:
      case ActionType::CANCEL: {
        if (order_rate_tracker_map_.count(symbol) == 0) {
          const auto& req_all = rate_limit_info_.req_all();
          order_rate_tracker_map_.emplace(
              symbol,
              std::make_unique<RateTracker>(req_all.secs() * 1'000'000'000LL));
          order_rate_tracker_map_[symbol]->SetRateLimitInfo(req_all.reqs(), req_all.reqs());
          ret = true;
        } else {
          ret = order_rate_tracker_map_[symbol]->HasEnoughRateLimit(ts);
        }
        break;
      }
      case ActionType::QUERY:
        // TODO(daniel):
        ret = true;
        break;
      default:
        CHECK(false) << "unsupported action: " << action_type;
        break;
    }
    if (!ret) {
      rate_limit_report_writer_->SetHitRateLimit(ts, action_type);
    }

    return ret;
  }

  void Push(int32_t num_reqs, int64_t ts, ActionType action_type, const std::string& symbol) {
    switch (action_type) {
      case ActionType::PLACE:
      case ActionType::AMEND:
      case ActionType::CANCEL: {
        if (order_rate_tracker_map_.count(symbol)) {
          order_rate_tracker_map_[symbol]->Push(num_reqs, ts);
        }
        break;
      }
      case ActionType::QUERY:
        // TODO(daniel):
        break;
      default:
        CHECK(false) << "unsupported action: " << action_type;
        break;
    }
  }

  void UpdateReqRateInfo(const nlohmann::json& symbol_rate_counts) {
    for (auto item : symbol_rate_counts.items()) {
      const auto symbol = item.key();
      if (order_rate_tracker_map_.count(symbol) == 0) {
        continue;
      }
      // TODO(daniel): just throw exception when ratecount is abnormal
      auto& rate_tracker = order_rate_tracker_map_[symbol];
      const auto rate_limit = rate_tracker->MaxReq();
      const auto rate_count = item.value().get<int>();
      const auto rate_limit_remaining = rate_limit - rate_count;
      rate_tracker->SetRateLimitInfo(rate_limit, rate_limit_remaining);
    }
  }

 private:
  std::unique_ptr<RateLimitReportWriter> rate_limit_report_writer_;
  coin2::exchange::order::RateLimitInfoSpotKraken rate_limit_info_;
  std::unordered_map<std::string, std::unique_ptr<RateTracker>> order_rate_tracker_map_;
};
